import warnings
import numpy as np

from ....utils.code_utils import validate_string

from ....survey import BaseSurvey
from ....utils import shift_to_discrete_topography
from . import receivers as Rx
from . import sources as Src
from ..utils import static_utils


class Survey(BaseSurvey):
    """DC/IP survey class.

    Parameters
    ----------
    source_list : list of simpeg.electromagnetic.static.resistivity.sources.BaseSrc
        List of SimPEG DC/IP sources
    survey_geometry : {"surface", "borehole", "general"}
        Survey geometry.
    """

    def __init__(
        self,
        source_list,
        survey_geometry="surface",
        **kwargs,
    ):
        if kwargs.pop("survey_type", None) is not None:
            raise TypeError(
                "Argument 'survey_type' has been removed in SimPEG 0.24.0. Types of sources and"
                "their corresponding receivers are obtained from their respective classes, without "
                "the need to specify the survey type.",
            )
        super(Survey, self).__init__(source_list, **kwargs)
        self.survey_geometry = survey_geometry

    @property
    def survey_geometry(self):
        """Survey geometry.

        Returns
        -------
        str
            Survey geometry; one of {"surface", "borehole", "general"}
        """
        return self._survey_geometry

    @survey_geometry.setter
    def survey_geometry(self, var):
        var = validate_string(
            "survey_geometry", var, ("surface", "borehole", "general")
        )
        self._survey_geometry = var

    @property
    def survey_type(self):
        """
        ``survey_type`` has been removed.

        .. important:

            The `survey_type` property has been removed. Types of sources and
            their corresponding receivers are obtained from their respective
            classes, without the need to specify the survey type.
        """
        raise AttributeError("'survey_type' has been removed.")

    @survey_type.setter
    def survey_type(self, var):
        raise AttributeError("'survey_type' has been removed.")

    def __repr__(self):
        return f"{self.__class__.__name__}(#sources: {self.nSrc}; #data: {self.nD})"

    @property
    def locations_a(self):
        """
        Locations of the positive (+) current electrodes in the survey.

        Returns
        -------
        (nD, dim) numpy.ndarray
            Locations of the positive (+) current electrodes in the survey
        """
        if getattr(self, "_locations_a", None) is None:
            self._set_abmn_locations()
        return self._locations_a

    @property
    def locations_b(self):
        """
        Locations of the negative (-) current electrodes in the survey.

        Returns
        -------
        (nD, dim) numpy.ndarray
            Locations of the negative (-) current electrodes in the survey
        """
        if getattr(self, "_locations_b", None) is None:
            self._set_abmn_locations()
        return self._locations_b

    @property
    def locations_m(self):
        """
        Locations of the positive (+) potential electrodes in the survey.

        Returns
        -------
        (nD, dim) numpy.ndarray
            Locations of the positive (+) potential electrodes in the survey
        """
        if getattr(self, "_locations_m", None) is None:
            self._set_abmn_locations()
        return self._locations_m

    @property
    def locations_n(self):
        """
        Locations of the negative (-) potential electrodes in the survey.

        Returns
        -------
        (nD, dim) numpy.ndarray
            Locations of the negative (-) potential electrodes in the survey
        """
        if getattr(self, "_locations_n", None) is None:
            self._set_abmn_locations()
        return self._locations_n

    @property
    def unique_electrode_locations(self):
        """
        Unique locations of the A, B, M, N electrodes.

        Returns
        -------
        (n, dim) numpy.ndarray
            The set of unique electrode locations used in the survey
        """
        loc_a = self.locations_a
        loc_b = self.locations_b
        loc_m = self.locations_m
        loc_n = self.locations_n
        return np.unique(np.vstack((loc_a, loc_b, loc_m, loc_n)), axis=0)

    @property
    def source_locations(self):
        """
        Returns, in order, the source locations for all sources in the survey.

        Returns
        -------
        list of numpy.ndarray
            List of length 2 containing the A and B electrode locations, in order.
        """
        src_a = []
        src_b = []

        for src in self.source_list:
            src_a.append(src.location_a)
            src_b.append(src.location_b)

        return [np.vstack(src_a), np.vstack(src_b)]

    def set_geometric_factor(
        self,
        space_type="halfspace",
    ):
        """
        Set and return the geometric factor for all data.

        Parameters
        ----------
        space_type : {'halfspace', 'wholespace'}
            Calculate geometric factors using a half-space or whole-space formula.

        Returns
        -------
        (nD) numpy.ndarray
            The geometric factor for each datum
        """
        geometric_factor = static_utils.geometric_factor(self, space_type=space_type)

        # geometric_factor = data.Data(self, geometric_factor)
        survey_slices = self.get_all_slices()
        for source in self.source_list:
            for rx in source.receiver_list:
                if rx.data_type == "apparent_resistivity":
                    src_rx_slice = survey_slices[source, rx]
                    rx._geometric_factor[source] = geometric_factor[src_rx_slice]
        return geometric_factor

    def _set_abmn_locations(self):
        locations_a = []
        locations_b = []
        locations_m = []
        locations_n = []
        for source in self.source_list:
            for rx in source.receiver_list:
                nRx = rx.nD
                # Pole Source
                if isinstance(source, Src.Pole):
                    locations_a.append(
                        source.location.reshape([1, -1]).repeat(nRx, axis=0)
                    )
                    locations_b.append(
                        source.location.reshape([1, -1]).repeat(nRx, axis=0)
                    )
                # Dipole Source
                elif isinstance(source, Src.Dipole):
                    locations_a.append(
                        source.location[0].reshape([1, -1]).repeat(nRx, axis=0)
                    )
                    locations_b.append(
                        source.location[1].reshape([1, -1]).repeat(nRx, axis=0)
                    )
                elif isinstance(source, Src.Multipole):
                    location_tiled = np.tile(source.location, (nRx, 1))
                    locations_a.append(location_tiled)
                    locations_b.append(location_tiled)
                # Pole RX
                if isinstance(rx, Rx.Pole):
                    locations_m.append(rx.locations)
                    locations_n.append(rx.locations)

                # Dipole RX
                elif isinstance(rx, Rx.Dipole):
                    locations_m.append(rx.locations[0])
                    locations_n.append(rx.locations[1])

        self._locations_a = np.vstack(locations_a)
        self._locations_b = np.vstack(locations_b)
        self._locations_m = np.vstack(locations_m)
        self._locations_n = np.vstack(locations_n)

    def drape_electrodes_on_topography(
        self,
        mesh,
        active_cells,
        topo_cell_cutoff="top",
        shift_horizontal=True,
        option=None,
        **kwargs,
    ):
        """Shift electrode locations to discrete surface topography.

        Parameters
        ----------
        mesh : discretize.TensorMesh or discretize.TreeMesh
            The mesh on which the discretized fields are computed
        active_cells : numpy.ndarray of int or bool
            Active topography cells
        topo_cell_cutoff : {"top", "center"}
            Define topography at tops of cells or cell centers.
        topography : (n, dim) numpy.ndarray, default = ``None``
            Surface topography

            .. deprecated:: v0.25.0

                The ``topography`` argument is not used in this function. It will be
                removed in SimPEG v0.27.0.

        force : bool, default = ``False``
            If ``True`` force electrodes to surface even if borehole

            .. deprecated:: v0.25.0

                The ``force`` argument is not used in this function. It will be removed
                in SimPEG v0.27.0.
        shift_horizontal : bool
            When True, locations are shifted horizontally to lie vertically over cell
            centers. When False, the original horizontal locations are preserved.
        option : {"top", "center"}
            Define topography at tops of cells or cell centers.

            .. deprecated:: 0.25.0

               Argument ``option`` is deprecated in favor of ``topo_cell_cutoff``
               and will be removed in SimPEG v0.27.0.

        """
        if option is not None:
            msg = (
                "Argument ``option`` is deprecated in favor of ``topo_cell_cutoff`` "
                "and will be removed in SimPEG v0.27.0."
            )
            warnings.warn(msg, FutureWarning, stacklevel=2)
            topo_cell_cutoff = option

        if (key := "topography") in kwargs:
            msg = (
                "The `topography` argument is not used in the "
                "`drape_electrodes_on_topography` and will be removed in "
                "SimPEG v0.27.0."
            )
            warnings.warn(msg, category=FutureWarning, stacklevel=2)
            kwargs.pop(key)

        if (key := "force") in kwargs:
            msg = (
                "The `force` argument is not used in the "
                "`drape_electrodes_on_topography` and will be removed in "
                "SimPEG v0.27.0."
            )
            warnings.warn(msg, category=FutureWarning, stacklevel=2)
            kwargs.pop(key)

        if kwargs:  # TODO Remove this when removing kwargs argument.
            raise TypeError("Unsupported keyword argument " + kwargs.popitem()[0])

        if self.survey_geometry == "surface":
            loc_a = self.locations_a[:, :2]
            loc_b = self.locations_b[:, :2]
            loc_m = self.locations_m[:, :2]
            loc_n = self.locations_n[:, :2]
            unique_electrodes, inv = np.unique(
                np.vstack((loc_a, loc_b, loc_m, loc_n)), return_inverse=True, axis=0
            )
            inv_a, inv = inv[: len(loc_a)], inv[len(loc_a) :]
            inv_b, inv = inv[: len(loc_b)], inv[len(loc_b) :]
            inv_m, inv_n = inv[: len(loc_m)], inv[len(loc_m) :]

            electrodes_shifted = shift_to_discrete_topography(
                mesh,
                unique_electrodes,
                active_cells=active_cells,
                topo_cell_cutoff=topo_cell_cutoff,
                shift_horizontal=shift_horizontal,
            )
            a_shifted = electrodes_shifted[inv_a]
            b_shifted = electrodes_shifted[inv_b]
            m_shifted = electrodes_shifted[inv_m]
            n_shifted = electrodes_shifted[inv_n]
            # These should all be in the same order as the survey datas
            ind = 0
            for src in self.source_list:
                a_loc, b_loc = a_shifted[ind], b_shifted[ind]
                if isinstance(src, (Src.Pole, Src.Multipole)):
                    src.location = [a_loc]
                else:
                    src.location = [a_loc, b_loc]
                for rx in src.receiver_list:
                    end = ind + rx.nD
                    m_locs, n_locs = m_shifted[ind:end], n_shifted[ind:end]
                    if isinstance(rx, Rx.Pole):
                        rx.locations = m_locs
                    else:
                        rx.locations = [m_locs, n_locs]
                    ind = end
            self._locations_a = a_shifted
            self._locations_b = b_shifted
            self._locations_m = m_shifted
            self._locations_n = n_shifted

        elif self.survey_geometry == "borehole":
            raise Exception("Not implemented yet for borehole survey_geometry")
        else:
            raise Exception(
                f"Input valid survey survey_geometry: {self.survey_geometry}"
            )
